#pragma once

namespace AsyncLib
{

class CAsyncMutex
{
public:
	CAsyncMutex() : m_nCount(0)
	{
	}

private:
	template<bool>
	friend class CAutoLock;

	static constexpr UINT Unique = 0x10000;

	auto LockAsync(bool bShared)
	{
		UINT nCount = bShared ? 1 : Unique;
		std::unique_lock lock(m_mutex);

		if (!m_pHead && TryAddCount(nCount))
			return CompletedTask();
		else
			return AddPendingTask(nCount);
	}

	bool TryLock(bool bShared)
	{
		UINT nCount = bShared ? 1 : Unique;
		std::unique_lock lock(m_mutex);
		return !m_pHead && TryAddCount(nCount);
	}

	void Unlock(bool bShared)
	{
		UINT nCount = bShared ? 1 : Unique;
		std::unique_lock lock(m_mutex);
		assert(m_nCount >= nCount);
		m_nCount -= nCount;

		while (m_pHead && TryAddCount(m_pHead->m_nCount))
		{
			m_pHead->GetState()->SetValue();
			Pop();
		}
	}

	bool TryAddCount(UINT nCount)
	{
		if (m_nCount + nCount <= Unique)
		{
			m_nCount += nCount;
			return true;
		}
		else
			return false;
	}

	class CStateManagerP : public CStateManager<void>
	{
	public:
		CStateManagerP(UINT nCount) : m_nCount(nCount)
		{
		}

		std::unique_ptr<CStateManagerP> m_pNext;
		const UINT m_nCount;
	};

	CCoTask<void> AddPendingTask(UINT nCount)
	{
		auto pState = std::make_shared<CAssocStateN<void, sizeof(CStateManagerP)>>();
		std::unique_ptr<CStateManagerP> pManager(new (pState->GetStorage()) CStateManagerP(nCount));
		pManager->SetState(pState);
		Push(std::move(pManager));
		return CCoTask<void>(std::move(pState));
	}

	void Push(std::unique_ptr<CStateManagerP>&& pNode)
	{
		auto ppNext = &m_pHead;
		for (; *ppNext; ppNext = &(*ppNext)->m_pNext);
		*ppNext = std::move(pNode);
	}

	void Pop()
	{
		m_pHead = std::move(m_pHead->m_pNext);
	}

	std::mutex m_mutex;
	std::unique_ptr<CStateManagerP> m_pHead; // Head of queue
	UINT m_nCount;
};

template<bool bShared>
class CAutoLock
{
public:
	CAutoLock(CAsyncMutex& mutex) : m_mutex(mutex), m_bLocked(false)
	{
	}

	~CAutoLock()
	{
		if (m_bLocked)
			m_mutex.Unlock(bShared);
	}

	CAutoLock(const CAutoLock&) = delete;
	CAutoLock& operator=(const CAutoLock&) = delete;

	[[nodiscard]] auto LockAsync()
	{
		assert(!m_bLocked);

		if (m_bLocked)
			Throw(std::errc::resource_deadlock_would_occur);

		m_bLocked = true;
		return m_mutex.LockAsync(bShared);
	}

	[[nodiscard]] bool TryLock()
	{
		assert(!m_bLocked);

		if (m_bLocked)
			Throw(std::errc::resource_deadlock_would_occur);

		return m_bLocked = m_mutex.TryLock(bShared);
	}

	void Unlock()
	{
		assert(m_bLocked);

		if (!m_bLocked)
			Throw(std::errc::operation_not_permitted);

		m_bLocked = false;
		m_mutex.Unlock(bShared);
	}

private:
	void Throw(std::errc ec)
	{
		throw std::system_error(std::make_error_code(ec));
	}

	CAsyncMutex& m_mutex;
	bool m_bLocked;
};

using CSharedLock = CAutoLock<true>;
using CUniqueLock = CAutoLock<false>;

} // namespace AsyncLib

namespace AL
{

using AsyncLib::CAsyncMutex;
using AsyncLib::CSharedLock;
using AsyncLib::CUniqueLock;

}
